#include "crypto_helper.hpp"
#include "hypercube/hypercube_cryto_agent.hpp"
#include "hypercube/hypercube_hmm_engine.hpp"
#include "mem_usage.hpp"
#include "mpi.h"
#include "mpi/mpi_util.hpp"
#include <sstream>

//* true or false
#define FIX_NTL_RAND_SEED(enable)                                              \
  do {                                                                         \
    NTL::SetSeed(NTL::conv<NTL::ZZ>(enable ? 0l : std::time(nullptr)));        \
  } while (0)

#define RAND_NTL_MATRIX(mat, row, col, bl)                                     \
  do {                                                                         \
    mat.SetDims(row, col);                                                     \
    for (long i = 0; i < row; i++) {                                           \
      for (long j = 0; j < col; j++) {                                         \
        mat[i][j] = NTL::RandomBits_ZZ(bl);                                    \
      }                                                                        \
    }                                                                          \
  } while (0)

int main(int argc, char **argv) {

  MPI_Init(&argc, &argv);
  bool sk;
  bool verbose;
  bool bsgs;
  bool single_thread;
  bool cache_free;
  bool minimal_m;

  unsigned long m, n, k;

  int rank, p;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &p);
  int sqrt_np = static_cast<int>(std::sqrt(p));
  p = sqrt_np * sqrt_np;
  // Defualt BGV Params Set
  params<helib::BGV> params;

  // Create a `command_line_parser` using `params<helib::BGV>`
  command_line_parser parser(params);

  // Add Custom Command Line Options
  parser.add_option("-m", "4", "dim-m", &m);
  parser.add_option("-n", "4", "dim-n", &n);
  parser.add_option("-k", "4", "dim-k", &k);
  parser.add_option("-v", "verbos", &verbose);
  parser.add_option("-sk", "using_secret_key", &sk);
  parser.add_option("-bsgs", "bsgs", &bsgs);
  parser.add_option("-st", "single_thread", &single_thread);
  parser.add_option("-nocache", "always_build_context", &cache_free);
  parser.add_option("-minimal", "find_minimal_m", &minimal_m);

  //* Parse command line
  int ret = parser.parse(argc, argv);
  if (ret) {
    return ret;
  }

  // Select encryption method according to the command line option `-sk`
  auto key = ENC::PUBKEY;

  // Create a HMM status instance
  hmm_status<shmm_engine> hmm_status;
  // Select a HMM method according to the command line option `-bsgs`
  hmm_status.method(bsgs ? METHOD::BSGS : METHOD::NAIVE);
  // Set # of threads according to the command line option `-st`
  if (!single_thread) {
    hmm_status.threads(omp_get_max_threads());
  }

  // Get the excepted hypercube size
  // hmme.status()->eh_cube();
  const auto eh_cube = hypercube_hmme::get_hypercube(m, n, k);

  // mem_begin_ = getCurrentRSS();
  // Initialize HE cryptograph system
  hypercube_hmmcc hmmcc(params, eh_cube, cache_free, minimal_m);
  // mem_end_ = getCurrentRSS();
  // mem_pre = static_cast<double>(mem_end_ - mem_begin_) / GB;

  std::stringstream ss;
  hmmcc.public_key().writeTo(ss);
  // mem_pk = static_cast<double>(ss.str().size()) / MB;

  // Print out the HElib context
  hmmcc.report_context();

  // Create an inactive HMM engine
  hypercube_hmme hmme(hmm_status);
  //* Register the HMM engine using EA
  const auto eh_mnk = hypercube_hmme::get_expected_mnk(m, n, k);
  hmme.register_engine(hmmcc.context()->getEA(), eh_mnk);
  std::string str_A, str_B;
  helib::Ctxt ctxt_A(hmmcc.public_key());
  helib::Ctxt ctxt_B(hmmcc.public_key());
  // Randomly generate input data
  NTL::mat_ZZ ptxt_A;
  NTL::mat_ZZ ptxt_B;
  if (rank == 0) {
    do {
      ptxt_A.SetDims(m, k);
      for (long i = 0; i < m; i++) {
        for (long j = 0; j < k; j++) {
          ptxt_A[i][j] = i * m + j;
        }
      }
    } while (0);
    do {
      ptxt_B.SetDims(k, n);
      for (long i = 0; i < k; i++) {
        for (long j = 0; j < n; j++) {
          ptxt_B[i][j] = i * k + j;
        }
      }
    } while (0);
    // FIX_NTL_RAND_SEED(false);
    //   RAND_NTL_MATRIX(ptxt_A, m, k, 8);
    //   RAND_NTL_MATRIX(ptxt_B, k, n, 8);

    // Encryption
    ctxt_A = hmmcc.encrypt(ptxt_A, key);
    ctxt_B = hmmcc.encrypt(ptxt_B, key);
  }

  if (rank == 0) {
    std::stringstream ssA, ssB;
    ctxt_A.writeTo(ssA);
    str_A = ssA.str();
    ctxt_B.writeTo(ssB);
    str_B = ssB.str();
  }

  MPI_Bcast_String(str_A, 0, MPI_COMM_WORLD);
  MPI_Bcast_String(str_B, 0, MPI_COMM_WORLD);

  if (rank != 0) {
    std::stringstream ss_A, ss_B;
    ss_A << str_A;
    ctxt_A = helib::Ctxt::readFrom(ss_A, hmmcc.public_key());
    ss_B << str_B;
    ctxt_B = helib::Ctxt::readFrom(ss_B, hmmcc.public_key());
    std::cout << rank << "readFrom finished" << std::endl;
  }

  // Execute HMM
  std::string partial_str;
  helib::Ctxt partial(hmmcc.public_key());
  partial.DummyEncrypt(NTL::ZZX(0l));
  auto ctxt_C = hmmcc.get_empty_ctxt();
  auto type = shmm_engine::parse_matrix_type(eh_mnk);
  hmme.cgemm(type, ctxt_C, ctxt_A, ctxt_B);
  partial += ctxt_C;
  std::cout << rank << "cgemm finished" << std::endl;
  std::stringstream ss_sum;
  partial.writeTo(ss_sum);
  partial_str = ss_sum.str();
  // ctxt_C.writeTo(ss_C);
  // mem_ctxt_C = static_cast<double>(ss_C.str().size()) / MB;

  // Decryption
  if (rank != 0) {
    MPI_Send_String(partial_str, 0, 0, MPI_COMM_WORLD);
  } else {
    std::vector<NTL::mat_ZZ> res;
    res.resize(p);
    for (int i = 0; i < p; i++) {
      std::stringstream ss;
      std::string ctxt_str;
      if (i == 0)
        ctxt_str = partial_str;
      else
        MPI_Recv_String(ctxt_str, i, 0, MPI_COMM_WORLD);
      ss << ctxt_str;
      auto ctxt = helib::Ctxt::readFrom(ss, hmmcc.public_key());
      hmmcc.decrypt(res[i], ctxt);
    }
    auto AB = ptxt_A * ptxt_B;
    for (std::size_t k = 0; k < p; k++) {
      const std::size_t i = k / sqrt_np;
      const std::size_t j = k % sqrt_np;
      for (std::size_t ii = 0; ii < 2; ii++) {
        for (std::size_t jj = 0; jj < 2; jj++) {
          auto expect = AB[i * 2 + ii][j * 2 + jj];
          // auto expect = ptxt_A[i * 4 + ii] * ptxt_B[j * 4 + jj];
          auto actual = res[k][ii][jj];
          if (expect != actual)
            std::cout << expect << ", " << actual << " ";
        }
        std::cout << std::endl;
      }
    }
    // if (verbose) {
    //   std::cout << "Dec(Enc(A) * Enc(B)):\n" << C << std::endl;
    // }

    // // Plaintext Verison
    // if (verbose) {
    //   auto ptxt_AB = ptxt_A * ptxt_B;
    //   std::cout << "A * B:\n" << ptxt_AB << std::endl;
    // }
  }
  // Report HMM
  // std::cout << "===================== HMM Report ======================\n";
  // hmmcc.report_time();
  // std::cout << "* Encryption Time            : " << enc_time << " ms\n";
  // std::cout << "* Decryption Time            : " << dec_time << " ms\n";
  // // std::cout << "* Enc(A) Size                : " << mem_ctxt_A << "
  // MB\n";
  // // std::cout << "* Enc(B) Size                : " << mem_ctxt_B << "
  // MB\n";
  // // std::cout << "* Enc(C) Size                : " << mem_ctxt_C << "
  // MB\n"; std::cout << "* Context Mem Usage          : " << mem_pre << "
  // GB\n"; std::cout << "  * Public Key Size          : " << mem_pk << "
  // MB\n"; std::cout << "  * Other(pre-computation)   : " << mem_pre -
  // (mem_pk / 1024.0)
  //           << " GB\n";
  // std::cout << "* HMM Mem Usage              : " << mem_algor << " MB\n";
  // std::cout << "* HMM Time (Outside)         : "
  //           << _named_local_timer_HMM.getTime() << " s\n";
  // std::cout << "=======================================================\n";
  // hmme.report();
  // std::cout << "=======================================================\n";
  MPI_Finalize();
  return 0;
}